{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #Checking Existence of Edge Length Limited Paths"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #union-find #graph #array #two-pointers #sorting"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #并查集 #图 #数组 #双指针 #排序"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: distanceLimitedPathsExist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #检查边长度限制的路径是否存在"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "<p>给你一个 <code>n</code> 个点组成的无向图边集 <code>edgeList</code> ，其中 <code>edgeList[i] = [u<sub>i</sub>, v<sub>i</sub>, dis<sub>i</sub>]</code> 表示点 <code>u<sub>i</sub></code> 和点 <code>v<sub>i</sub></code> 之间有一条长度为 <code>dis<sub>i</sub></code> 的边。请注意，两个点之间可能有 <strong>超过一条边 </strong>。</p>\n",
    "\n",
    "<p>给你一个查询数组<code>queries</code> ，其中 <code>queries[j] = [p<sub>j</sub>, q<sub>j</sub>, limit<sub>j</sub>]</code> ，你的任务是对于每个查询 <code>queries[j]</code> ，判断是否存在从 <code>p<sub>j</sub></code> 到 <code>q<sub>j</sub></code><sub> </sub>的路径，且这条路径上的每一条边都 <strong>严格小于</strong> <code>limit<sub>j</sub></code> 。</p>\n",
    "\n",
    "<p>请你返回一个 <b>布尔数组</b><em> </em><code>answer</code><em> </em>，其中<em> </em><code>answer.length == queries.length</code> ，当 <code>queries[j]</code> 的查询结果为 <code>true</code> 时， <code>answer</code> 第<em> </em><code>j</code> 个值为<em> </em><code>true</code><em> </em>，否则为 <code>false</code> 。</p>\n",
    "\n",
    "<p> </p>\n",
    "\n",
    "<p><strong>示例 1：</strong></p>\n",
    "<img alt=\"\" src=\"https://assets.leetcode-cn.com/aliyun-lc-upload/uploads/2020/12/19/h.png\" style=\"width: 267px; height: 262px;\" />\n",
    "<pre>\n",
    "<b>输入：</b>n = 3, edgeList = [[0,1,2],[1,2,4],[2,0,8],[1,0,16]], queries = [[0,1,2],[0,2,5]]\n",
    "<b>输出：</b>[false,true]\n",
    "<b>解释：</b>上图为给定的输入数据。注意到 0 和 1 之间有两条重边，分别为 2 和 16 。\n",
    "对于第一个查询，0 和 1 之间没有小于 2 的边，所以我们返回 false 。\n",
    "对于第二个查询，有一条路径（0 -> 1 -> 2）两条边都小于 5 ，所以这个查询我们返回 true 。\n",
    "</pre>\n",
    "\n",
    "<p><strong>示例 2：</strong></p>\n",
    "<img alt=\"\" src=\"https://assets.leetcode-cn.com/aliyun-lc-upload/uploads/2020/12/19/q.png\" style=\"width: 390px; height: 358px;\" />\n",
    "<pre>\n",
    "<b>输入：</b>n = 5, edgeList = [[0,1,10],[1,2,5],[2,3,9],[3,4,13]], queries = [[0,4,14],[1,4,13]]\n",
    "<b>输出：</b>[true,false]\n",
    "<b>解释：</b>上图为给定数据。\n",
    "</pre>\n",
    "\n",
    "<p> </p>\n",
    "\n",
    "<p><strong>提示：</strong></p>\n",
    "\n",
    "<ul>\n",
    "\t<li><code>2 <= n <= 10<sup>5</sup></code></li>\n",
    "\t<li><code>1 <= edgeList.length, queries.length <= 10<sup>5</sup></code></li>\n",
    "\t<li><code>edgeList[i].length == 3</code></li>\n",
    "\t<li><code>queries[j].length == 3</code></li>\n",
    "\t<li><code>0 <= u<sub>i</sub>, v<sub>i</sub>, p<sub>j</sub>, q<sub>j</sub> <= n - 1</code></li>\n",
    "\t<li><code>u<sub>i</sub> != v<sub>i</sub></code></li>\n",
    "\t<li><code>p<sub>j</sub> != q<sub>j</sub></code></li>\n",
    "\t<li><code>1 <= dis<sub>i</sub>, limit<sub>j</sub> <= 10<sup>9</sup></code></li>\n",
    "\t<li>两个点之间可能有 <strong>多条</strong> 边。</li>\n",
    "</ul>\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [checking-existence-of-edge-length-limited-paths](https://leetcode.cn/problems/checking-existence-of-edge-length-limited-paths/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [checking-existence-of-edge-length-limited-paths](https://leetcode.cn/problems/checking-existence-of-edge-length-limited-paths/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['3\\n[[0,1,2],[1,2,4],[2,0,8],[1,0,16]]\\n[[0,1,2],[0,2,5]]', '5\\n[[0,1,10],[1,2,5],[2,3,9],[3,4,13]]\\n[[0,4,14],[1,4,13]]']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UF:\n",
    "    def __init__(self, n):\n",
    "        self.rank = [1] * n\n",
    "        self.parent = list(range(n))\n",
    "        self.timestamp = [inf] * n\n",
    "\n",
    "    def union(self, u, v, time):\n",
    "        p1 = self.find(u)\n",
    "        p2 = self.find(v)\n",
    "        if p1 != p2:\n",
    "            if self.rank[p1] > self.rank[p2]:\n",
    "                self.timestamp[p2] = time\n",
    "                self.parent[p2] = p1\n",
    "                self.rank[p1] += self.rank[p2]\n",
    "            else:\n",
    "                self.timestamp[p1] = time\n",
    "                self.parent[p1] = p2\n",
    "                self.rank[p2] += self.rank[p1]\n",
    "            \n",
    "    def find(self, x, time = inf):\n",
    "        if x != self.parent[x] and self.timestamp[x] < time:\n",
    "            x = self.find(self.parent[x], time)\n",
    "        return x\n",
    "        \n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        uf, res = UF(n), []\n",
    "        edgeList.sort(key=lambda x: x[2])\n",
    "        for x, y, time in edgeList:\n",
    "            uf.union(x, y, time)\n",
    "        for p, q, limit in queries:\n",
    "            res.append(uf.find(p, limit) == uf.find(q, limit))\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        def find(x):\n",
    "            if p[x]!=x:\n",
    "                p[x]=find(p[x])\n",
    "            return p[x]\n",
    "        p=list(range(n))\n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "\n",
    "        j=0\n",
    "        ans=[False]*len(queries)\n",
    "        for i ,(a,b,limit) in sorted(enumerate(queries), key=lambda x: x[1][2]):\n",
    "            while j<len(edgeList) and edgeList[j][2]<limit:\n",
    "                u,v,_=edgeList[j]\n",
    "                p[find(u)]=find(v)\n",
    "                j+=1\n",
    "            ans[i]=(find(a)==find(b))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int) -> None:\n",
    "        self.root = [i for i in range(n)]\n",
    "        self.size = [1] * n\n",
    "        self.part = n\n",
    "        return\n",
    "\n",
    "    def find(self, x):\n",
    "        lst = []\n",
    "        while x != self.root[x]:\n",
    "            lst.append(x)\n",
    "            # 在查询的时候合并到顺带直接根节点\n",
    "            x = self.root[x]\n",
    "        for w in lst:\n",
    "            self.root[w] = x\n",
    "        return x\n",
    "\n",
    "    def union(self, x, y):\n",
    "        root_x = self.find(x)\n",
    "        root_y = self.find(y)\n",
    "        if root_x == root_y:\n",
    "            return False\n",
    "        if self.size[root_x] >= self.size[root_y]:\n",
    "            root_x, root_y = root_y, root_x\n",
    "        self.root[root_x] = root_y\n",
    "        self.size[root_y] += self.size[root_x]\n",
    "        # 将非根节点的秩赋0\n",
    "        self.size[root_x] = 0\n",
    "        self.part -= 1\n",
    "        return True\n",
    "\n",
    "    def is_connected(self, x, y):\n",
    "        return self.find(x) == self.find(y)\n",
    "\n",
    "    def get_root_part(self):\n",
    "        # 获取每个根节点对应的组\n",
    "        part = defaultdict(list)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            part[self.find(i)].append(i)\n",
    "        return part\n",
    "\n",
    "    def get_root_size(self):\n",
    "        # 获取每个根节点对应的组大小\n",
    "        size = defaultdict(int)\n",
    "        n = len(self.root)\n",
    "        for i in range(n):\n",
    "            size[self.find(i)] = self.size[self.find(i)]\n",
    "        return size\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "\n",
    "        m = len(queries)\n",
    "        ans = [False]*m\n",
    "        ind = list(range(m))\n",
    "        ind.sort(key=lambda it: queries[it][2])\n",
    "        uf = UnionFind(n)\n",
    "        edgeList.sort(key=lambda it: it[2])\n",
    "        k = len(edgeList)\n",
    "        j = 0\n",
    "        for i in ind:\n",
    "            limit = queries[i][2]\n",
    "            while j < k and edgeList[j][2] < limit:\n",
    "                x, y = edgeList[j][:-1]\n",
    "                uf.union(x, y)\n",
    "                j += 1\n",
    "            a, b = queries[i][:-1]\n",
    "            if uf.is_connected(a, b):\n",
    "                ans[i] = True\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key=lambda e: e[2])\n",
    "\n",
    "        # 并查集模板\n",
    "        fa = list(range(n))\n",
    "        def find(x: int) -> int:\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def merge(from_: int, to: int) -> None:\n",
    "            fa[find(from_)] = find(to)\n",
    "\n",
    "        ans, k = [False] * len(queries), 0\n",
    "        # 查询的下标按照 limit 从小到大排序，方便离线\n",
    "        for i, (p, q, limit) in sorted(enumerate(queries), key=lambda p: p[1][2]):\n",
    "            while k < len(edgeList) and edgeList[k][2] < limit:\n",
    "                merge(edgeList[k][0], edgeList[k][1])\n",
    "                k += 1\n",
    "            ans[i] = find(p) == find(q)\n",
    "        return ans\n",
    "\n",
    "# 作者：力扣官方题解\n",
    "# 链接：https://leetcode.cn/problems/checking-existence-of-edge-length-limited-paths/solutions/2018397/jian-cha-bian-chang-du-xian-zhi-de-lu-ji-cdr5/\n",
    "# 来源：力扣（LeetCode）\n",
    "# 著作权归作者所有。商业转载请联系作者获得授权，非商业转载请注明出处。"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        f=list(range(n))\n",
    "        def find(x):\n",
    "            if f[x]==x:\n",
    "                return x\n",
    "            fx=find(f[x])\n",
    "            f[x]=fx\n",
    "            return fx\n",
    "        \n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "        j=0\n",
    "        res=[False]*len(queries)\n",
    "        for i, (a, b, limit) in sorted(enumerate(queries), key=lambda x:x[1][2]):\n",
    "            while j<len(edgeList) and edgeList[j][2]<limit:\n",
    "                u, v, _ = edgeList[j]\n",
    "                f[find(u)]=find(v)\n",
    "                j+=1\n",
    "            res[i]=find(a)==find(b)\n",
    "        \n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key = lambda x : x[2])\n",
    "        fa = [x for x in range(n)]\n",
    "        def find(x):\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        def union(x,y):\n",
    "            fa[find(x)] = find(y)\n",
    "        \n",
    "        ans = [False]*len(queries)\n",
    "        k = 0\n",
    "        for i,(p,q,limit) in sorted(enumerate(queries),key = lambda x: x[1][2]):\n",
    "            while k < len(edgeList) and edgeList[k][2]<limit:\n",
    "                union(edgeList[k][0],edgeList[k][1])\n",
    "                k+=1\n",
    "            ans[i] = find(p) == find(q)\n",
    "        \n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key=lambda a: a[2])\n",
    "\n",
    "        fa = list(range(n))\n",
    "\n",
    "        def find(x: int) -> int :\n",
    "            if x != fa[x]:\n",
    "               fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "\n",
    "        def union(from_ :int , to_ : int) -> None:\n",
    "            fa[find(from_)] = find(to_)\n",
    "        \n",
    "        ans, k = [False] * len(queries), 0\n",
    "\n",
    "        for i , (p,q,limit) in sorted(enumerate(queries),key=lambda p: p[1][2]):\n",
    "            while k < len(edgeList) and edgeList[k][2] < limit:\n",
    "                union(edgeList[k][0],edgeList[k][1])\n",
    "                k = k + 1\n",
    "            ans[i] = find(p) == find(q)\n",
    "\n",
    "        return ans\n",
    "        \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        ans = [False] * len(queries)\n",
    "        edgeList.sort(key=lambda x: x[2])\n",
    "        queries = sorted(enumerate(queries), key=lambda x: x[1][2])\n",
    "\n",
    "        UF = UnionFind(n)\n",
    "        j = 0\n",
    "        m = len(edgeList)\n",
    "        for i, (u, v, limit) in queries:\n",
    "            while j < m and edgeList[j][2] < limit:\n",
    "                UF.unionSet(edgeList[j][0], edgeList[j][1])\n",
    "                j += 1\n",
    "            ans[i] = UF.isSameSet(u, v)\n",
    "        \n",
    "        return ans \n",
    "\n",
    "        \n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n: int, m: int=None) -> None:\n",
    "        \"\"\"\n",
    "        build m disjoint sets from n nodes\n",
    "        T: O(N)\n",
    "        \"\"\"\n",
    "        self.parents = list(range(n))   \n",
    "        self.ranks = [0] * n            \n",
    "    \n",
    "    def findSet(self, i: int) -> int:\n",
    "        \"\"\"return root of node i's set\n",
    "           T: amortized O(α(N))  effectively O(1)\n",
    "             first O(N) then O(1)\n",
    "           path-compression  \n",
    "           Recursion\n",
    "        \"\"\"\n",
    "        if self.parents[i] == i:\n",
    "            return i\n",
    "        else:\n",
    "            self.parents[i] = self.findSet(self.parents[i])\n",
    "            return self.parents[i]\n",
    "        \n",
    "\n",
    "    \n",
    "    def unionSet(self, i: int, j: int) -> None:\n",
    "        \"\"\"若节点i和节点j来自2个不同的互斥集合，将两个集合合并 \n",
    "           union by rank   矮树作为高树的子树\n",
    "           T: amortized O(α(N))  effectively O(1)\n",
    "             first O(N) then O(1)\n",
    "        \"\"\"\n",
    "        ip, jp = self.findSet(i), self.findSet(j)\n",
    "\n",
    "        # 若两个节点来自同一互斥集合，返回\n",
    "        if ip == jp:\n",
    "            return \n",
    "\n",
    "        # 若A树矮，B树高，A树连到B树的根节点上，新树rank不变\n",
    "        if self.ranks[ip] < self.ranks[jp]:\n",
    "            self.parents[ip] = jp\n",
    "\n",
    "        # 若B树矮，A树高，B树连到A树的根节点上，新树rank不变\n",
    "        elif self.ranks[jp] < self.ranks[ip]:\n",
    "            self.parents[jp] = ip\n",
    "\n",
    "        # 若两树等高，任选一种方案，新树 rank +1\n",
    "        else:\n",
    "            self.parents[jp] = ip\n",
    "            self.ranks[ip] += 1\n",
    "    \n",
    "    def isSameSet(self, i: int, j: int) -> bool:\n",
    "        \"\"\"判断节点i和节点j是否属于同一集合\n",
    "           T: O(α(N) first O(N) then O(1)\n",
    "        \"\"\"\n",
    "        # path compression twice\n",
    "        return self.findSet(i) == self.findSet(j)\n",
    "    \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        e=sorted(edgeList,key=lambda x:x[2])\n",
    "        #print(e)\n",
    "        fa=[ i for i in range(n)]\n",
    "        def find(x):\n",
    "            if fa[x]==x:\n",
    "                return x\n",
    "            fa[x]=find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(a,b):\n",
    "            fa[find(a)]=find(b)\n",
    "        m=len(queries)\n",
    "        ans=[False]*m\n",
    "        k=0\n",
    "        for i,(a,b,l) in sorted(enumerate(queries),key= lambda x:x[1][2]):\n",
    "            #print(i,a,b,l)\n",
    "            while k<len(edgeList) and e[k][2]<l:\n",
    "                #print(e[k])\n",
    "                if find(e[k][0])!=find(e[k][1]):\n",
    "                    union(e[k][0],e[k][1])\n",
    "                k+=1\n",
    "            #print(fa)\n",
    "            ans[i]=(find(a)==find(b))\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UF:\n",
    "    def __init__(self, n):\n",
    "        self.rank = [1] * n\n",
    "        self.parent = list(range(n))\n",
    "\n",
    "    def union(self, n1, n2):\n",
    "        p1 = self.find(n1)\n",
    "        p2 = self.find(n2)\n",
    "        if self.rank[p1] > self.rank[p2]:\n",
    "            self.parent[p2] = p1\n",
    "            self.rank[p1] += self.rank[p2]\n",
    "        else:\n",
    "            self.parent[p1] = p2\n",
    "            self.rank[p2] += self.rank[p1]\n",
    "            \n",
    "    def find(self, n):\n",
    "        while n != self.parent[n]:\n",
    "            n = self.parent[n]\n",
    "        return n\n",
    "        \n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key=lambda x: x[2])\n",
    "        uf, res, k = UF(n), [False] * len(queries), 0\n",
    "        for i, (p, q, limit) in sorted(enumerate(queries), key=lambda p: p[1][2]):\n",
    "            while k < len(edgeList) and edgeList[k][2] < limit:\n",
    "                uf.union(edgeList[k][0], edgeList[k][1])\n",
    "                k += 1\n",
    "            res[i] = uf.find(p) == uf.find(q)\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key = lambda p: p[2])\n",
    "        fa = list(range(n))\n",
    "        def find(x:int) -> int:\n",
    "            if x != fa[x]:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(x:int, y:int):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if fx != fy:\n",
    "                fa[fx] = fy\n",
    "        ans = [False] * len(queries)\n",
    "        j = 0\n",
    "        for i, (a, b , limit) in sorted(enumerate(queries), key = lambda p: p[1][2]):\n",
    "            while j < len(edgeList) and edgeList[j][2] < limit:\n",
    "                u, v, _ = edgeList[j]\n",
    "                union(u, v)\n",
    "                j += 1\n",
    "            ans[i] = find(a) == find(b)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from typing import List\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parent = list(range(n))\n",
    "        \n",
    "    def find(self, u):\n",
    "        if self.parent[u] == u:\n",
    "            return u\n",
    "        self.parent[u] = self.find(self.parent[u])\n",
    "        return self.parent[u]\n",
    "    \n",
    "    def merge(self, u, v):\n",
    "        u = self.find(u)\n",
    "        v = self.find(v)\n",
    "        if u == v:\n",
    "            return\n",
    "        self.parent[u] = v\n",
    "        \n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        edgeList.sort(key=lambda x: x[2])\n",
    "        queries = sorted(enumerate(queries), key=lambda x: x[1][2])\n",
    "        \n",
    "        uf = UnionFind(n)\n",
    "        answer = [False] * len(queries)\n",
    "        edge_idx = 0\n",
    "        \n",
    "        for i, (p, q, limit) in queries:\n",
    "            while edge_idx < len(edgeList) and edgeList[edge_idx][2] < limit:\n",
    "                uf.merge(edgeList[edge_idx][0], edgeList[edge_idx][1])\n",
    "                edge_idx += 1\n",
    "            \n",
    "            answer[i] = (uf.find(p) == uf.find(q))\n",
    "        \n",
    "        return answer\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "pa = [0] * 100000\r\n",
    "rk = [0] * 100000\r\n",
    "\r\n",
    "def find(i: int) -> int:\r\n",
    "    if i != pa[i]: pa[i] = find(pa[i])\r\n",
    "    return pa[i]\r\n",
    "\r\n",
    "def join(i: int, j: int) -> None:\r\n",
    "    f = find(i)\r\n",
    "    g = find(j)\r\n",
    "    if f == g: return\r\n",
    "    if rk[f] > rk[g]:\r\n",
    "        pa[g] = f\r\n",
    "    else:\r\n",
    "        if rk[f] == rk[g]:\r\n",
    "            rk[g] += 1\r\n",
    "        pa[f] = g\r\n",
    "    return\r\n",
    "\r\n",
    "class Solution:\r\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\r\n",
    "        for i in range(n):\r\n",
    "            pa[i] = i\r\n",
    "            rk[i] = 0\r\n",
    "        \r\n",
    "        pnt = 0\r\n",
    "        m = len(edgeList)\r\n",
    "        edgeList.sort(key=lambda e: e[2])\r\n",
    "        res = [False] * len(queries)\r\n",
    "\r\n",
    "        for i, (u, v, lmt) in sorted(enumerate(queries), key=lambda e: e[1][2]):\r\n",
    "            while pnt < m and edgeList[pnt][2] < lmt:\r\n",
    "                join(edgeList[pnt][0], edgeList[pnt][1])\r\n",
    "                pnt += 1\r\n",
    "            res[i] = find(u) == find(v)\r\n",
    "        \r\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        qs = sorted((limit, p, q, i) for i, (p, q, limit) in enumerate(queries))\n",
    "        edgeList.sort(key=lambda x: x[2])\n",
    "        m = len(queries)\n",
    "        res = [False] * m\n",
    "        \n",
    "        fa = list(range(n + 1))\n",
    "        def find(x):\n",
    "            if fa[x] != x:\n",
    "                fa[x] = find(fa[x])\n",
    "            return fa[x]\n",
    "        def union(x, y):\n",
    "            fx, fy = find(x), find(y)\n",
    "            fa[fx] = fy\n",
    "        \n",
    "        i = 0\n",
    "        for limit, p, q, qi in qs:\n",
    "            while i < len(edgeList) and edgeList[i][2] < limit:\n",
    "                union(edgeList[i][0], edgeList[i][1])\n",
    "                i += 1\n",
    "            if find(p) == find(q):\n",
    "                res[qi] = True\n",
    "        return res"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parents = [i for i in range(n)]\n",
    "    \n",
    "    def find(self, u):\n",
    "        if u != self.parents[u]:\n",
    "            self.parents[u] = self.find(self.parents[u])\n",
    "        return self.parents[u]\n",
    "\n",
    "    def union(self, u, v):\n",
    "        pu, pv = self.find(u), self.find(v)\n",
    "        if pu != pv:\n",
    "            self.parents[pv] = pu\n",
    "    \n",
    "    def isConnected(self, u, v):\n",
    "        pu, pv = self.find(u), self.find(v)\n",
    "        return pu == pv\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        # 把query按照limit排序\n",
    "        # 先把小于limit的边加进并查集\n",
    "        # 然后看u, v是否相连\n",
    "        # 随着下一个query的limit增大，又有一些可用的边能加入并查集\n",
    "        # 依次遍历query, 每次并查集里的边都是小于当前query的limit的\n",
    "        uf = UnionFind(n)\n",
    "        edgeList.sort(key=itemgetter(2))\n",
    "        for i in range(len(queries)):\n",
    "            queries[i].append(i)\n",
    "        queries.sort(key=itemgetter(2))\n",
    "        \n",
    "        res = [0] * len(queries)\n",
    "        i = 0\n",
    "        for u, v, limit, idx in queries:\n",
    "            while i < len(edgeList) and edgeList[i][2] < limit:\n",
    "                uf.union(edgeList[i][0], edgeList[i][1])\n",
    "                i += 1\n",
    "            res[idx] = uf.isConnected(u, v)\n",
    "\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        def find(x):\n",
    "            if x != father[x]:\n",
    "                father[x] = find(father[x])\n",
    "            return father[x]\n",
    "        def isSameSet(x,y):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if fx != fy:\n",
    "                return False\n",
    "            return True\n",
    "        def union(x,y):\n",
    "            fx = find(x)\n",
    "            fy = find(y)\n",
    "            if fx != fy:\n",
    "                father[fx] = fy\n",
    "\n",
    "        m = len(edgeList)\n",
    "        father = [x for x in range(n)]\n",
    "        ans = [0]*len(queries)\n",
    "        for i in range(len(queries)):\n",
    "            queries[i].append(i)\n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "        queries.sort(key=lambda x:x[2])\n",
    "        k = 0\n",
    "        for p,q,l,i in queries:\n",
    "            while k<m and edgeList[k][2]<l:\n",
    "                union(edgeList[k][0],edgeList[k][1])\n",
    "                k += 1\n",
    "            ans[i] = isSameSet(p,q)\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self, n):\n",
    "        self.parents = [i for i in range(n)]\n",
    "    \n",
    "    def find(self, u):\n",
    "        if u != self.parents[u]:\n",
    "            self.parents[u] = self.find(self.parents[u])\n",
    "        return self.parents[u]\n",
    "\n",
    "    def union(self, u, v):\n",
    "        pu, pv = self.find(u), self.find(v)\n",
    "        if pu != pv:\n",
    "            self.parents[pv] = pu\n",
    "    \n",
    "    def isConnected(self, u, v):\n",
    "        pu, pv = self.find(u), self.find(v)\n",
    "        return pu == pv\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        uf = UnionFind(n)\n",
    "        edgeList.sort(key=itemgetter(2))\n",
    "        for i in range(len(queries)):\n",
    "            queries[i].append(i)\n",
    "        queries.sort(key=itemgetter(2))\n",
    "        \n",
    "        res = [0] * len(queries)\n",
    "        i = 0\n",
    "        for u, v, limit, idx in queries:\n",
    "            while i < len(edgeList) and edgeList[i][2] < limit:\n",
    "                uf.union(edgeList[i][0], edgeList[i][1])\n",
    "                i += 1\n",
    "            res[idx] = uf.isConnected(u, v)\n",
    "\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        fa=list(range(n))\n",
    "        def find(a):\n",
    "            if fa[a]!=a:\n",
    "                fa[a]=find(fa[a])\n",
    "            return fa[a]\n",
    "        edgeList.sort(key=lambda x:-x[2])\n",
    "        queries=[[i]+queries[i] for i in range(len(queries))]\n",
    "        res=[0]*len(queries)\n",
    "        queries.sort(key=lambda x:x[3])\n",
    "        for i,u,v,limit in queries:\n",
    "            while edgeList and edgeList[-1][2]<limit:\n",
    "                a,b,_=edgeList.pop()\n",
    "                fa[find(a)]=find(b)\n",
    "            res[i]= find(u)==find(v)\n",
    "        return res\n",
    "                "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class UnionFind:\n",
    "    def __init__(self,n):\n",
    "        self.parent = list(range(n))\n",
    "        self.size = [1]*n\n",
    "\n",
    "    def find(self,x):\n",
    "        while not x==self.parent[x]:\n",
    "            self.parent[x] = self.parent[self.parent[x]]\n",
    "            x = self.parent[x]\n",
    "        return x\n",
    "    \n",
    "    def union(self, x, y):\n",
    "        rootX = self.find(x)\n",
    "        rootY = self.find(y)\n",
    "        if rootX == rootY:\n",
    "            return\n",
    "        if self.size[rootX] > self.size[rootY]:\n",
    "            rootX,rootY = rootY, rootX\n",
    "        self.parent[rootX] = rootY\n",
    "        self.size[rootY] += self.size[rootX]\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        qlen = len(queries)\n",
    "        ans = [False]*qlen\n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "        queries = [[p,q,limit,i] for i,(p,q,limit) in enumerate(queries)]\n",
    "        queries.sort(key=lambda x:x[2])\n",
    "        uf = UnionFind(n)\n",
    "        pos_e, pos_q = 0,0\n",
    "        while pos_q<qlen:\n",
    "            p,q,limit,idx = queries[pos_q]\n",
    "            while pos_e<len(edgeList) and edgeList[pos_e][2] < queries[pos_q][2]:\n",
    "                u,v,d = edgeList[pos_e]\n",
    "                uf.union(u, v)\n",
    "                pos_e += 1\n",
    "            ans[idx] = uf.find(p)==uf.find(q)\n",
    "            pos_q+=1\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        m = len(edgeList)\n",
    "        k = len(queries)\n",
    "        questions = [[0,0,0,0] for _ in range(k)]\n",
    "        for i in range(k):\n",
    "            questions[i][0] = queries[i][0]\n",
    "            questions[i][1] = queries[i][1]\n",
    "            questions[i][2] = queries[i][2]\n",
    "            questions[i][3] = i\n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "        questions.sort(key=lambda x:x[2])\n",
    "        father = list(range(n))\n",
    "        def find(i):\n",
    "            if i != father[i]:\n",
    "                father[i] = find(father[i])\n",
    "            return father[i]\n",
    "        def union(x, y):\n",
    "            father[find(x)] = find(y)\n",
    "        def issame(x, y):\n",
    "            return find(x) == find(y)\n",
    "        ans = [False for _ in range(k)]\n",
    "        j = 0\n",
    "        for i in range(k):\n",
    "            while j < m and edgeList[j][2] < questions[i][2]:\n",
    "                union(edgeList[j][0], edgeList[j][1])\n",
    "                j += 1\n",
    "            ans[questions[i][3]] = issame(questions[i][0], questions[i][1])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        m = len(edgeList)\n",
    "        k = len(queries)\n",
    "        questions = [[0,0,0,0] for _ in range(k)]\n",
    "        for i in range(k):\n",
    "            questions[i][0] = queries[i][0]\n",
    "            questions[i][1] = queries[i][1]\n",
    "            questions[i][2] = queries[i][2]\n",
    "            questions[i][3] = i\n",
    "        edgeList.sort(key=lambda x:x[2])\n",
    "        questions.sort(key=lambda x:x[2])\n",
    "        father = list(range(n))\n",
    "        def find(i):\n",
    "            if i != father[i]:\n",
    "                father[i] = find(father[i])\n",
    "            return father[i]\n",
    "        def union(x, y):\n",
    "            father[find(x)] = find(y)\n",
    "        def issame(x, y):\n",
    "            return find(x) == find(y)\n",
    "        ans = [False for _ in range(k)]\n",
    "        j = 0\n",
    "        for i in range(k):\n",
    "            while j < m and edgeList[j][2] < questions[i][2]:\n",
    "                union(edgeList[j][0], edgeList[j][1])\n",
    "                j += 1\n",
    "            ans[questions[i][3]] = issame(questions[i][0], questions[i][1])\n",
    "        return ans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def distanceLimitedPathsExist(self, n: int, edgeList: List[List[int]], queries: List[List[int]]) -> List[bool]:\n",
    "        parent = {i : i for i in range(n)}\n",
    "        def find(node):\n",
    "            path = []\n",
    "            while node != parent[node]:\n",
    "                path.append(node)\n",
    "                node = parent[node]\n",
    "            for p in path:\n",
    "                parent[p] = node \n",
    "            return node\n",
    "        def union(n1, n2):\n",
    "            p1 = find(n1)\n",
    "            p2 = find(n2)\n",
    "            if p1 != p2:\n",
    "                parent[p1] = p2 \n",
    "                \n",
    "        qr = sorted([(idx, q[0], q[1], q[2]) for idx, q in enumerate(queries)], key = lambda x: x[-1])\n",
    "        edgeList.sort(key = lambda x: - x[-1])\n",
    "        res = [False] * len(queries)\n",
    "        for idx, p, q, limit in qr:\n",
    "            while edgeList and edgeList[-1][-1] < limit:\n",
    "                a, b, _ = edgeList.pop()\n",
    "                union(a, b)\n",
    "            if find(p) == find(q):\n",
    "                res[idx] = True\n",
    "        return res\n",
    "            "
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
